import asyncio
import aiohttp
import aiofiles
from pathlib import Path
import json

# BASE_URL = "http://localhost:8016/v1/ocr/stepFun"
BASE_URL = "http://localhost:18002/v1/ocr/stepFun"
TEST_IMAGE_FILE = 'image_ocr.jpg'


async def test_upload_file(session):
    """测试图片上传"""
    print("\n[测试] POST /v1/orc/stepFun")

    # 准备图片文件上传
    async with aiofiles.open(TEST_IMAGE_FILE, 'rb') as f:
        file_content = await f.read()

    # 创建 FormData
    data = aiohttp.FormData()
    data.add_field('file',
                   file_content,
                   filename=TEST_IMAGE_FILE,
                   content_type='image/jpeg')
    data.add_field('process', 'false')
    data.add_field('model', 'StepFun/GOT-OCR-2.0-hf')

    async with session.post(f"{BASE_URL}", data=data) as response:
        print(f'response.status = {response.status =}')
        assert response.status == 200
        result = await response.json()
        print(f"  响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
        assert result["status"] == "uploaded"
        assert result["filename"] == TEST_IMAGE_FILE.name
        print("  ✓ 图片上传测试通过")
        return result


async def run_all_tests():
    """运行所有测试"""
    print("=" * 60)
    print("开始运行测试套件")
    print("=" * 60)

    async with aiohttp.ClientSession() as session:
        try:
            # 测试文件上传
            upload_result = await test_upload_file(session)
            uploaded_filename = upload_result["filename"]

            print("\n" + "=" * 60)
            print("所有测试完成！")
            print("=" * 60)

        except AssertionError as e:
            pass
            # print(f"\n❌ 测试失败: {e}")
            # raise
        except aiohttp.ClientError as e:
            pass
            # print(f"\n❌ 连接错误: {e}")
            # print("请确保 FastAPI 服务器正在运行 (python main.py)")
            # raise
        except Exception as e:
            pass
            # print(f"\n❌ 发生错误: {e}")
            # raise
        finally:
            # 清理测试文件
            if TEST_IMAGE_FILE.exists():
                TEST_IMAGE_FILE.unlink()
                print(f"\n清理: 删除测试文件 {TEST_IMAGE_FILE}")

if __name__ == "__main__":
    print("使用 aiohttp 测试 FastAPI 音频上传 API")
    print(f"目标服务器: {BASE_URL}")
    print("\n提示: 请确保服务器正在运行 (python main.py)")

    try:
        asyncio.run(run_all_tests())
    except KeyboardInterrupt:
        print("\n\n测试被用户中断")
    except Exception as e:
        print(f"\n测试失败: {e}")

